import os
import cv2
import numpy as np
from queue import Queue, LifoQueue, PriorityQueue

imageDir = os.path.abspath("images")
filenames = os.listdir(imageDir)

# k为提取色数
def medianSegmentation(imgPath, k = 256):
	# 读取图片	
	img = cv2.imdecode(np.fromfile(imgPath, dtype=np.uint8), -1)
	img = np.array(img)
	# 3维转2维
	img = img.reshape((-1, img.shape[2]))
	# 队列
	que = Queue()
	# 入队
	que.put(img)
	# 切分数量到k之前不停止
	while que.qsize() < k:
		# 取色块
		img = que.get()
		# 找rgb中的最大轴
		selectColor = 0
		selectDifference = 0
		for i in range(3):
			selectMinVal = np.min(img[:, i])
			selectMaxVal = np.max(img[:, i])
			if selectMaxVal - selectMinVal > selectDifference:
				selectColor = i
				selectDifference = selectMaxVal - selectMinVal
		# 使用最大轴进行排序
		sortedIndexs = np.lexsort((img[:, selectColor], ))
		img = img[sortedIndexs, :]
		# 切分
		leftImg = img[:img.shape[0]//2, :]
		rightImg = img[img.shape[0]//2:, :]
		# 入队列
		que.put(leftImg)
		que.put(rightImg)
	# 提取颜色
	colors = []
	while not que.empty():
		# 取色块
		img = que.get()
		# 取像素均值
		colors.append(np.mean(img, axis=0))
	# 显示色块
	showColors = np.zeros(shape=(len(colors) * 20, 200, 3), dtype=np.uint8)
	for i in range(k):
		showColors[i*20: i*20+20, :, :] = colors[i]
	cv2.imshow("", showColors)
	cv2.waitKey(0)

def main():
	medianSegmentation(os.path.join(imageDir, filenames[0]), 16)

if "__main__" == __name__:
	main()
	exit(0)